import random

import numpy as np
import pandas as pd
import torch
from matplotlib import pyplot as plt
from scipy.io import arff
from sklearn.metrics import roc_curve, auc, roc_auc_score
from sklearn.model_selection import train_test_split
from torch import nn, device
from torch.utils.data import Dataset, WeightedRandomSampler
from torchvision.transforms import transforms
import warnings

# 忽略所有警告
warnings.filterwarnings("ignore")

# 检查GPU是否可用
if torch.cuda.is_available():
    # 将模型移动到GPU上
    device = 'cuda'
else:
    device = 'cpu'


class CNN(nn.Module):  # cnn
    def __init__(self):
        super(CNN, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv1d(in_channels=60,out_channels=32,
                      kernel_size=3,stride=1,padding=1),
            nn.ReLU()
        )
        self.layer2 = nn.Sequential(
            nn.Conv1d(in_channels=16, out_channels=32,
                      kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
        )
        self.layer3 = nn.Sequential(
            nn.Conv1d(in_channels=16, out_channels=8, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
        )
        self.max = nn.MaxPool1d(kernel_size=2, stride=2)
        self.drop_out = nn.Dropout()
        self.fc1 = nn.Linear(4, 32)
        self.fc2 = nn.Linear(32, 2)

    def forward(self, x):
        out = self.layer1(x)
        out = out.T
        out = self.max(out)
        out = out.T
        out = self.layer2(out)
        out = out.T
        out = self.max(out)
        out = out.T
        out = self.layer3(out)
        out = out.T
        out = self.max(out)
        out = self.drop_out(out)
        out = self.fc1(out)
        out = self.fc2(out)

        return out




class getDataSet(Dataset):
    def __init__(self, train=True, train_x=None, train_y=None, test_x=None, test_y=None, val=False, transform=None):
        self.train = train
        self.val = val
        self.transform = transform
        if self.train:
            self.dataset = train_x
            self.labels = train_y
        elif val:
            self.dataset = test_x
            self.labels = test_y
        else:
            self.dataset = test_x

    def __getitem__(self, index):
        if self.train:
            return torch.Tensor(self.dataset[index]).to(device), self.labels[index].to(device)
        elif self.val:
            return torch.Tensor(self.dataset[index]).to(device), self.labels[index].to(device)
        else:
            return torch.Tensor(self.dataset[index]).to(device)

    def __len__(self):
        return self.dataset.shape[0]

# 计算正反例比例
def get_ratio(labels):
    # 计算正类别和负类别的样本数量
    num_positive_samples = (labels == 1).sum().item()
    num_negative_samples = (labels == 0).sum().item()

    # 计算每个样本的权重，使得正类别和负类别的样本被采样的概率相等
    weights = [1.0 / num_positive_samples if label == 1 else 1.0 / num_negative_samples for label in labels]

    # 创建 WeightedRandomSampler
    sampler = WeightedRandomSampler(weights, len(weights), replacement=True)
    return sampler

def main():
    # 处理数据
    # file_name = 'dataSet/AEEEM/JDT.arff'
    # data, meta = arff.loadarff(file_name)
    # file_name = 'datasets/JDT.csv'
    file_name = 'feature/JDT.csv'
    # 读取csv
    data = pd.read_csv(file_name, header=None)
    data = pd.DataFrame(data)
    data = data.drop([0])


    datasets = data.iloc[:, :-1]
    labels = data.iloc[:, -1]

    # 将特征数据转为数组
    datasets = np.array(datasets)
    # print(datasets)
    # 标签的转换为0、1
    labels = np.array(labels)

    for i in range(len(labels)):
        if labels[i] == "b'clean'":
            labels[i] = np.int32(0)
        else:
            labels[i] = np.int32(1)

    num = datasets.shape[0]
    over = num - (num // 81)  * 81

    # 随机删除数据
    for i in range(over):
        index = random.randint(0, num - (i+1))
        datasets = np.delete(datasets, index, axis=0)
        labels = np.delete(labels, index, axis=0)

    # 训练集和测试集的划分
    x_train, x_test, y_train, y_test = train_test_split(datasets, labels, test_size=0.25)
    sampler = get_ratio(y_train)
    sampler1 = get_ratio(y_test)
    # 转为tensor
    x_train = torch.Tensor(x_train.astype(float))
    x_test = torch.Tensor(x_test.astype(float))
    y_train = torch.Tensor(y_train.astype(int))
    y_test = torch.Tensor(y_test.astype(int))

    # 特征数据转置
    # x_train = x_train.T
    # x_test = x_test.T

    model = CNN().to(device)  # 模型

    learning_rate = 0.001  # 学习率
    # criterion = nn.MSELoss()  # 损失函数 均方误差
    # criterion = nn.L1Loss()  # 损失函数 平均绝对误差
    # criterion = nn.SmoothL1Loss()  # 损失函数 平滑L1损失
    criterion = nn.CrossEntropyLoss()  # 损失函数 交叉熵损失
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)  # 优化器
    batch_size = 81  # 批处理大小
    num_epochs = 5  # 训练次数
    # 定义dataset
    train_dataset = getDataSet(train_x=x_train, train_y=y_train, transform=transforms.ToTensor())
    test_dataset = getDataSet(train=False, test_x=x_test, test_y=y_test, val=True, transform=transforms.ToTensor())
    # 定义loader
    train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                               batch_size=batch_size,
                                               sampler=sampler,
                                               shuffle=False)

    test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                              batch_size=batch_size,
                                              sampler=sampler1,
                                              shuffle=False)

    total_step = len(train_loader)
    loss_list = []
    acc_list = []
    recall_list = []
    for epoch in range(num_epochs):
        for i, (data, labels) in enumerate(train_loader):
            data = data.T
            # 前向传播
            outputs = model(data)
            # sm = nn.Softmax(dim=1)
            # outputs = sm(outputs)
            loss = criterion(outputs, labels.long())
            loss_list.append(loss.item())

            # 反向传播，执行Adam优化
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            # 记录准确率
            total = labels.size(0)
            predicted = torch.argmax(outputs.data, dim=1)
            correct = (predicted == labels).sum().item()
            acc_list.append(correct / total)



            # 召回率
            TP = torch.sum((predicted == 1) & (labels == 1))
            FN = torch.sum((predicted == 0) & (labels == 1))
            recall = TP / (TP + FN)
            # 转为numpy
            recall = recall.cpu().numpy()
            recall_list.append(recall)
            labels = labels.cpu().numpy()
            predicted = predicted.cpu().numpy()
            # ROC曲线
            fpr, tpr, thresholds = roc_curve(labels, predicted)
            roc_auc = auc(fpr, tpr)
            # AUC值
            auc_score = roc_auc_score(labels, predicted)

            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}, Accuracy: {:.2f}%, Recall: {:.2f}%, AUC: {:.2f}%'
                  .format(epoch + 1, num_epochs, i + 1, total_step, loss.item(),
                          (correct / total) * 100, recall * 100, auc_score * 100))
            # 选择准确率、召回率和AUC均较好的模型
            if correct / total > 0.7 and recall > 0.7 and auc_score > 0.7:
                # torch.save(model.state_dict(), 'model/cnnM/model_{:.3f}.ckpt'.format(correct / total))
                # 保存模型，格式.pth
                torch.save(model.state_dict(), 'model/cnnM/model_{:.3f}.pth'.format(correct / total))

            # 根据训练次数画出准确率、召回率和AUC的变化曲线



    # 测试
    model.eval()
    with torch.no_grad():
        correct = 0
        total = 0
        recall = 0
        for data, labels in test_loader:
            data = data.to(device)
            data =data.T
            labels = labels.to(device)
            outputs = model(data)
            # sm = nn.Softmax(dim=1)
            # outputs = sm(outputs)
            # 记录准确率
            total += labels.size(0)
            predicted = torch.argmax(outputs.data, dim=1)
            correct += (predicted == labels).sum().item()
            acc_list.append(correct / total)

            # 召回率
            TP = torch.sum((predicted == 1) & (labels == 1))
            FN = torch.sum((predicted == 0) & (labels == 1))
            recall += TP / (TP + FN)
            # 转为numpy
            recall = recall.cpu().numpy()
            recall_list.append(recall)
            labels = labels.cpu().numpy()
            predicted = predicted.cpu().numpy()
            # ROC曲线
            fpr, tpr, thresholds = roc_curve(labels, predicted)
            roc_auc = auc(fpr, tpr)
            # AUC值
            auc_score = roc_auc_score(labels, predicted)

            print('Test Accuracy of the model on the test datas: {} %'.format((correct / total) * 100))

            plt.figure()
            plt.plot(acc_list, color='darkorange', lw=2, label='Acc curve')
            plt.xlim([-0.01, 9.0])
            plt.ylim([0.0, 1.01])
            plt.xlabel('Step')
            plt.ylabel('Acc')
            plt.title('Acc example')
            plt.legend(loc="lower right")
            plt.show()

            plt.figure()
            plt.plot(recall_list, color='darkorange', lw=2, label='Recall curve')
            plt.xlim([-0.01, 9.0])
            plt.ylim([0.0, 1.01])
            plt.xlabel('Step')
            plt.ylabel('Recall')
            plt.title('Recall example')
            plt.legend(loc="lower right")
            plt.show()

            plt.figure()
            plt.plot(fpr, tpr, color='darkorange', lw=2, label='ROC curve')
            plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
            plt.xlim([-0.01, 1.0])
            plt.ylim([0.0, 1.01])
            plt.xlabel('False Positive Rate')
            plt.ylabel('True Positive Rate')
            plt.title('Receiver operating characteristic example')
            plt.legend(loc="lower right")
            plt.show()

if __name__ == '__main__':
    main()